package com.zjj.lbw.ai.component;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.aggregator.DefaultContentAggregator;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.rag.query.router.LanguageModelQueryRouter;
import dev.langchain4j.rag.query.transformer.ExpandingQueryTransformer;
import dev.langchain4j.service.TokenStream;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import opennlp.tools.tokenize.lang.en.TokenSampleStream;

import java.util.Collection;
import java.util.List;
import java.util.Map;

public class QueryTransformTest {
    public static void main(String[] args) {
        ChatLanguageModel model = OpenAiChatModel.builder().modelName("gpt-4o-mini").apiKey("demo").build();

        // 向量存储数据库
        EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        // 向量模型
        OpenAiEmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .baseUrl("http://langchain4j.dev/demo/openai/v1")
                .apiKey("demo")
                .modelName("text-embedding-3-small")
                .build();

        ContentRetriever contentRetriever1 = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(3) // 最相似的3个结果
                .minScore(0.7) // 只找相似度在0.8以上的内容
                .build();
        ContentRetriever contentRetriever2 = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(3) // 最相似的3个结果
                .minScore(0.7) // 只找相似度在0.8以上的内容
                .build();

        Query query = Query.from("余额什么时候到账");
        // ExpandingQueryTransformer 可以根据原始的查询内容，扩展出三个类似语义的查询内容，然后可以根据扩展内容 进行查询出三组内容
        ExpandingQueryTransformer transformer = new ExpandingQueryTransformer(model);
        Collection<Query> result = transformer.transform(query);

        // 向量匹配 q1、q2、q2      1，2，3  1，4，5  1，3，6
        // 内容增强器，将查询到的向量结果进行增强，默认的是按出现的次数，算分聚合
        DefaultContentAggregator contentAggregator = new DefaultContentAggregator();

        // 查询路由器，如果知识库很大、支持的场景很多，可以分成不同的知识库
        // 然后根据不同的查询内容，路由到不同知识库，ContentRetriever 内容检索器 就可以理解为不同的向量库
        LanguageModelQueryRouter queryRouter = new LanguageModelQueryRouter(model,
                Map.of(contentRetriever1, "订单知识库", contentRetriever2, "商品知识库"));
        Collection<ContentRetriever> retrievers = queryRouter.route(Query.from("今天新增了那些商品"));

        // TokenStream 是指的用 AiService这种代理方式，支持流式响应，代理的方法返回值应该声明为 TokenStream这种
        // onNext() 方法处理正常响应，onError()处理错误响应，start()方法开启
    }
}
